import logging  # 用于日志记录
import time     # 用于时间相关操作
from ruamel import yaml  # 用于YAML格式配置文件的读写（ruamel.yaml支持更完善的YAML语法）


class AverageMeter(object):
    """
    指标统计工具类，用于计算和存储数值的当前值、平均值、总和及计数
    常用于训练过程中记录损失、准确率等指标的实时统计
    """
    def __init__(self):
        # 初始化时重置所有统计值
        self.reset()

    def reset(self):
        """重置所有统计参数（通常在每个epoch或迭代开始时调用）"""
        self.val = 0    # 当前值（最新一次更新的数值）
        self.avg = 0    # 平均值（sum / count）
        self.sum = 0    # 总和（所有值的累加和）
        self.count = 0  # 计数（参与统计的样本数量）

    def update(self, val, n=1):
        """
        更新统计值
        :param val: 新的数值（如当前batch的损失值）
        :param n: 该数值对应的样本数量（默认1，如batch_size为32时n=32）
        """
        self.val = val          # 更新当前值
        self.sum += val * n     # 累加总和（数值×样本数，确保平均计算准确）
        self.count += n         # 累加样本计数
        self.avg = self.sum / self.count  # 重新计算平均值


def Prepare_logger(args, eval=False):
    """
    配置并返回日志记录器（logger），同时输出到控制台和文件
    :param args: 命令行参数对象，需包含snapshot_pref（日志保存路径前缀）
    :param eval: 是否为评估模式（决定日志文件名）
    :return: 配置好的logger对象
    """
    # 获取名为当前模块的logger实例
    logger = logging.getLogger(__name__)
    # 关闭日志传播（避免被父logger重复处理）
    logger.propagate = False
    # 设置日志级别为INFO（只记录INFO及以上级别日志）
    logger.setLevel(logging.INFO)

    # 创建控制台处理器（输出日志到控制台）
    handler = logging.StreamHandler()
    # 定义日志格式（包含时间、级别、消息）
    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
    handler.setFormatter(formatter)
    handler.setLevel(0)  # 控制台输出所有级别的日志
    logger.addHandler(handler)

    # 生成带时间戳的日志文件名（避免重复）
    date = time.strftime('%Y%m%d%H%M', time.localtime(time.time()))  # 格式：年月日时分
    # 训练模式日志路径：snapshot_pref/时间.log；评估模式：snapshot_pref/时间-Eval.log
    logfile = args.snapshot_pref + date + '.log' if not eval else args.snapshot_pref + f'/{date}-Eval.log'
    
    # 创建文件处理器（输出日志到文件）
    file_handler = logging.FileHandler(logfile, mode='w')  # 'w'表示覆盖写入
    file_handler.setLevel(logging.INFO)  # 文件只记录INFO及以上级别
    file_handler.setFormatter(formatter)  # 使用相同的日志格式
    logger.addHandler(file_handler)

    return logger


def get_configs(dataset):
    """
    从数据集配置文件中获取指定数据集的配置
    :param dataset: 数据集名称（用于索引配置）
    :return: 该数据集的配置字典（如路径、类别数等）
    """
    # 加载数据集配置文件（./configs/dataset_cfg.yaml）
    data = yaml.load(open('./configs/dataset_cfg.yaml'))
    # 返回指定数据集的配置
    return data[dataset]


def get_and_save_args(parser):
    """
    解析命令行参数，与默认配置合并，并保存最终配置到文件
    功能：将命令行输入的参数覆盖默认配置，生成当前实验的完整配置
    :param parser: 已定义的argparse解析器
    :return: 合并后的最终配置字典
    """
    # 解析命令行参数，得到参数对象
    args = parser.parse_args()

    # 加载默认配置文件（./configs/default_config.yaml）
    # 使用RoundTripLoader保留原文件格式（如注释、顺序）
    default_config = yaml.load(open('./configs/default_config.yaml', 'r'), Loader=yaml.RoundTripLoader)
    # 将命令行参数转换为字典（便于键值对操作）
    current_config = vars(args)

    # 合并配置：用命令行参数覆盖默认配置（仅覆盖已指定的参数）
    for k, v in current_config.items():
        # 若命令行参数在默认配置中存在
        if k in default_config:
            # 若命令行参数不为空且与默认值不同，则更新
            if (v != default_config[k]) and (v is not None):
                print(f"Updating:  {k}: {default_config[k]} (default) ----> {v}")
                default_config[k] = v

    # 将合并后的配置保存到当前目录（./current_configs.yaml）
    # 使用RoundTripDumper保留格式，便于人类阅读
    yaml.dump(default_config, open('./current_configs.yaml', 'w'), indent=4, Dumper=yaml.RoundTripDumper)
    
    # 返回最终配置（用于实验中使用）
    return default_config